{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac8d8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1b11b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def im_show(image):\n",
    "    plt.imshow(image)\n",
    "    plt.show()\n",
    "    \n",
    "def read_view_in(in_path):\n",
    "    frame_id_2_persons = {}\n",
    "    with open(in_path, 'r') as fin:\n",
    "        for line in fin.readlines():\n",
    "            items = line.split(',')\n",
    "            frame_id = int(items[0])\n",
    "            if frame_id in frame_id_2_persons:\n",
    "                frame_id_2_persons[frame_id].append(items[1:])\n",
    "            else:\n",
    "                frame_id_2_persons[frame_id] = [items[1:]]\n",
    "    return frame_id_2_persons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02dab5f0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#input_video = '/mnt/nas/04_自由视角处理后素材/球员识别测试164/view_004_test.ts'\n",
    "\n",
    "# input_video = 'view_out/029_to_004_4.mp4'\n",
    "# input_view_poin_path = 'view_out/034_to_004.txt'\n",
    "# output_video = 'view_out/034_to_004_5.mp4'\n",
    "\n",
    "input_video = 'view_out/034_to_004_5.mp4'\n",
    "input_view_poin_path = 'boxes/view_004_boxes.txt'\n",
    "output_video = 'view_out/0004_to_004_6_1.mp4'\n",
    "\n",
    "frame_id_2_persons = read_view_in(input_view_poin_path)\n",
    "frame_id = -1\n",
    "more = True\n",
    "\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "#width, height = 1920, 1080\n",
    "print(width, height)\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_id += 1\n",
    "    if(frame_id % 20 == 0):\n",
    "        print(\"== handle \", frame_id)\n",
    "    if(frame_id in frame_id_2_persons):\n",
    "        persons = frame_id_2_persons[frame_id]\n",
    "        for person in persons:\n",
    "            num_txt, x, y , width, height = person\n",
    "            x,y,height = int(x), int(y),int(height)\n",
    "            y = y+height\n",
    "            if(x>0 and y>0):\n",
    "                cv2.circle(frame, (x,y), 3, (255, 255, 255), 5)\n",
    "                #cv2.putText(frame, num_txt, (x,y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)\n",
    "    writer.write(frame)\n",
    "    \n",
    "reader.release()\n",
    "writer.release()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dec3030",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "def read_valid_number(num_path):\n",
    "    valid_num = []\n",
    "    with open(num_path, 'r') as fin:\n",
    "        for line in fin.readlines():\n",
    "            valid_num.append(line.strip())\n",
    "    return valid_num\n",
    "\n",
    "valid_nums = read_valid_number('./valid_num.txt')\n",
    "\n",
    "input_view_point_pathes = ['view_out/010_to_004.txt', \n",
    "                           'view_out/016_to_004.txt','view_out/023_to_004.txt', \n",
    "                           'view_out/029_to_004.txt', 'view_out/034_to_004.txt']\n",
    "\n",
    "boxes_path = 'boxes/view_004_boxes.txt'\n",
    "#input_video = 'view_out/0004_to_004_6.mp4'\n",
    "input_video = '/mnt/nas/04_自由视角处理后素材/球员识别测试164/view_004_test.ts'\n",
    "output_video = \"view_out/final_demo4.mp4\"\n",
    "\n",
    "frame_id = -1\n",
    "more = True\n",
    "frame_id_2_person_list = list(map(read_view_in, input_view_point_pathes))\n",
    "\n",
    "frame_id_2_person_box = read_view_in(boxes_path)\n",
    "\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]\n",
    "\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_id += 1\n",
    "    if(frame_id % 20 == 0):\n",
    "        print(\"== handle \", frame_id)\n",
    "        \n",
    "    person_2_boxes = {}\n",
    "    for idx, frame_id_2_persons in enumerate(frame_id_2_person_list):\n",
    "        if(frame_id in frame_id_2_persons):\n",
    "            #print(\"== frame id \", frame_id)\n",
    "            persons = frame_id_2_persons[frame_id]\n",
    "            for person in persons:\n",
    "                num_txt, x, y, p_width, p_height = person\n",
    "                if num_txt not in valid_nums:\n",
    "                    print(\"== ignore invalid num \", num_txt)\n",
    "                    continue\n",
    "                x, y, p_width, p_height = int(x), int(y), int(p_width), int(p_height)\n",
    "                if(x>0 and x<width and y>0 and y<height):\n",
    "                    if(num_txt in person_2_boxes):\n",
    "                        person_2_boxes[num_txt].append((idx, x, y, p_width, p_height))\n",
    "                    else:\n",
    "                        person_2_boxes[num_txt] = [(idx, x, y, p_width, p_height)]\n",
    "    \n",
    "    num_2_box_xy = {}\n",
    "    for person_num in person_2_boxes:\n",
    "        persons = person_2_boxes[person_num]\n",
    "        x = 0\n",
    "        y = 0\n",
    "        c_y = 0\n",
    "        for person in persons:\n",
    "            idx, p_x, p_y, p_width, p_height = person\n",
    "            color = colors[idx]\n",
    "            #cv2.circle(frame, (x,y), 3, color, 5)\n",
    "            \n",
    "            x += p_x\n",
    "            y += p_y-p_height\n",
    "            c_y += p_y\n",
    "        x = int(x/len(persons))\n",
    "        y = int(y/len(persons))\n",
    "        c_y = int(c_y/len(persons))\n",
    "        num_2_box_xy[person_num] = (x, y, c_y)\n",
    "        \n",
    "#     for num_txt in num_2_box_xy:\n",
    "#         x, y, c_y = num_2_box_xy[num_txt]\n",
    "#         cv2.putText(frame, num_txt, (x,y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 1)\n",
    "#         cv2.circle(frame, (x,c_y), 3, (0, 0, 0), 5)\n",
    "        \n",
    "    final_num_2_box = num_2_box_xy\n",
    "    if frame_id in frame_id_2_person_box:\n",
    "        view_0_persons = frame_id_2_person_box[frame_id]\n",
    "        for person in view_0_persons:\n",
    "            num_txt, p_x, p_y, p_width, p_height = person\n",
    "            if(num_txt not in valid_nums):\n",
    "                print(\"== ignore invalid view 0 num \", num_txt)\n",
    "                continue\n",
    "            p_x, p_y, p_width, p_height = int(p_x), int(p_y), int(p_width), int(p_height)\n",
    "            p_c_y = p_y+p_height\n",
    "            \n",
    "            if num_txt in num_2_box_xy:\n",
    "                x, y, c_y = num_2_box_xy[num_txt]\n",
    "                x += p_x\n",
    "                y += p_y\n",
    "                c_y += p_c_y\n",
    "                x = int(x/2)\n",
    "                y = int(y/2)\n",
    "                c_y = int(c_y/2)\n",
    "                final_num_2_box[num_txt] = (x, y, c_y)\n",
    "            else:\n",
    "                final_num_2_box[num_txt] = (p_x, p_y, p_c_y)\n",
    "                \n",
    "    for num_txt in final_num_2_box:\n",
    "        x, y, c_y = final_num_2_box[num_txt]\n",
    "        cv2.putText(frame, num_txt, (x,y), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 1)\n",
    "        cv2.circle(frame, (x,c_y), 3, (255, 255, 255), 5)\n",
    "    \n",
    "    writer.write(frame)\n",
    "    \n",
    "reader.release()\n",
    "writer.release()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "272148ff",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
